
package br.unifor.g2cl;

import java.io.IOException;
import java.net.SocketAddress;
import java.util.HashMap;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;

import net.sf.jgcs.Annotation;
import net.sf.jgcs.ClosedSessionException;
import net.sf.jgcs.ControlSession;
import net.sf.jgcs.ExceptionListener;
import net.sf.jgcs.GroupConfiguration;
import net.sf.jgcs.JGCSException;
import net.sf.jgcs.Message;
import net.sf.jgcs.MessageListener;
import net.sf.jgcs.NotJoinedException;
import net.sf.jgcs.Service;
import net.sf.jgcs.ServiceListener;
import net.sf.jgcs.UnsupportedServiceException;
import net.sf.jgcs.membership.MembershipSession;

/* TODO
 * Importante fazer essa classe verificar a saída de membros
 */

/**
 * DataSession that fragments messages larger than the fragSize.
 * The message is divided into smaller messages and re-assembled at
 * receiving side. 
 *           1b        4b           4b            4b         ....
 * header: header message ID | message size | frag offset | payload
 * @author henrique
 */
public class FragDataSession implements IMarshalDataSession, ExceptionListener,
		MessageListener, ServiceListener {

	private static Logger logger = Logger.getLogger(FragDataSession.class
			.getName());

	private static final byte HEADER = 0xA;

	
	// XXX synchronize it?
	private HashMap<SocketAddress, HashMap<Integer, AsmMessage>> fragments;

	private IMarshalDataSession data;

	private ExceptionListener realExceptionListener;
	private MessageListener realMessageListener;
	private ServiceListener realServiceListener;
	
	private MembershipSession membershipSession;

	// XXX must be configurable
	private int fragSize = 3;

	// XXX need to be atomic int ???
	//private AtomicInteger myNextMessageId = new AtomicInteger(0);
	private int myNextMessageId;

	public void setFragSize(String fragSize) {
		this.fragSize = Integer.parseInt(fragSize);
	}
	
	public int getNextMessageId() {
		//return myNextMessageId.incrementAndGet(); 
		return ++myNextMessageId;
	}

	public FragDataSession(IMarshalDataSession data, ControlSession control) {		
		this.data = data;
		this.fragments = new HashMap<SocketAddress, HashMap<Integer, AsmMessage>>();
		if (control instanceof MembershipSession) {
			this.membershipSession = (MembershipSession) control;
			// TODO add a listener
		}
	}

	public int getFragSize() {
		return fragSize;
	}

	public void close() {
		data.close();
	}

	public Message createMessage() throws ClosedSessionException {
		return data.createMessage();
	}

	public GroupConfiguration getGroup() {
		return data.getGroup();
	}

	public void multicast(Message msg, Service service, Object cookie,
			Annotation... annotation) throws IOException,
			UnsupportedServiceException {
		FragmentGenerator fg = new FragmentGenerator((G2CLMessage) msg, fragSize);
		while (fg.hasNext()) {
			data.multicast(fg.next(), service, cookie, annotation);
		}
		
	}

	public void send(Message msg, Service service, Object cookie,
			SocketAddress destination, Annotation... annotation)
			throws IOException, UnsupportedServiceException {
		FragmentGenerator fg = new FragmentGenerator((G2CLMessage) msg, fragSize);
		while (fg.hasNext()) {
			data.send(fg.next(), service, cookie, destination, annotation); 
		}

	}

	public void setExceptionListener(ExceptionListener exception)
			throws ClosedSessionException {
		this.realExceptionListener = exception;
		data.setExceptionListener(this);
	}

	public void setMessageListener(MessageListener listener)
			throws ClosedSessionException {
		this.realMessageListener = listener;
		data.setMessageListener(this);
	}

	public void setServiceListener(ServiceListener listener)
			throws ClosedSessionException {
		this.realServiceListener = listener;
		data.setServiceListener(this);
	}

	public Object onMessage(Message msg) {
		G2CLMessage message = (G2CLMessage) msg;
		
		byte[] header = message.removeHeader();
		
		if (header.length == 13 && header[0] == HEADER) {
			byte[] payload = msg.getPayload();
			SocketAddress sender = msg.getSenderAddress();
			
			int msgId = Util.byteArrayToInt(header, 1);
			int msgSize = Util.byteArrayToInt(header, 5);
			int msgOffset = Util.byteArrayToInt(header, 9);
			
			if (logger.isLoggable(Level.FINEST)) {
				logger.finest("Fragmented message received. msgId: " + msgId +
						", msgSize: " + msgSize + ", msgOffset: " + msgOffset);
			}
			
			if (!fragments.containsKey(sender)) {
				fragments.put(sender, new HashMap<Integer, AsmMessage>());
			}
			HashMap<Integer, AsmMessage> peerFragTable = fragments.get(sender);
			
			if (!peerFragTable.containsKey(msgId)) {
				peerFragTable.put(msgId, new AsmMessage(msgSize));
			}
			
			AsmMessage asm = peerFragTable.get(msgId);
			asm.add(payload, msgOffset);
			
			if (asm.isComplete()) {
				try {
					if (logger.isLoggable(Level.FINER)) {
						logger.finer("Message id " + msgId + " re-assembled. "+
								"Delivering it...");
					}
					G2CLMessage bigMessage = (G2CLMessage) data.createMessage();
					bigMessage.setPayload(asm.getMessage());
					peerFragTable.remove(msgId);
					bigMessage.setSenderAddress(sender);
					return realMessageListener.onMessage(bigMessage);
				} catch (ClosedSessionException e) {
					logger.log(Level.WARNING, "received message and then " +
							"disconnected? (the last message was not delivered)", e);
				}
			}
			
			
		} else {
			if (logger.isLoggable(Level.SEVERE)) {
			logger.severe("ERROR! Discarding message. This is not a " +
					"fragmented message. Make sure that the stack of " +
					"datasessions are configured same way at the two " +
					"endpoints.");
			}
		}
		
		return null;
	}

	public void onException(JGCSException exception) {
		realExceptionListener.onException(exception);
	}

	public void onServiceEnsured(Object context, Service service) {
		realServiceListener.onServiceEnsured(context, service);
	}
	
	
	public void onExcluded() {
		logger.finest("Excluded from group.");
		fragments.clear();			
	}

	public void onMembershipChange() {
		try {
			List<SocketAddress> failed = membershipSession.getMembership().getFailedMembers();
			List<SocketAddress> leaved = membershipSession.getMembership().getLeavedMembers();
			
			for (SocketAddress s : failed) {
				fragments.remove(s);
			}

			for (SocketAddress s : leaved) {
				fragments.remove(s);
			}
			
		} catch (NotJoinedException e) {
			logger.log(Level.WARNING, "received membership change and then " +
					"disconnected?", e);
		}
	}


	/** Iterator-like class to generate fragments of a message */
	 private class FragmentGenerator {

		private int fragSize;
		private byte[] payload;
		private int total;
		private int p;
		private G2CLMessage message;
		private int messageId;

		public FragmentGenerator(G2CLMessage message, int fragSize) {
			this.message = message;
			this.fragSize = fragSize;
			this.payload = message.getPayload();
			this.total =  payload.length;
			this.p = 0;
			this.messageId = getNextMessageId();
			
			if (total > fragSize && logger.isLoggable(Level.FINER)) {
				logger.finer("message id " + messageId + " will be fragmented");
			}
			
		}

		public boolean hasNext() {
			return p < total;
		}

		public G2CLMessage next() throws ClosedSessionException {
			// size of the next fragment
			int size = Math.min(fragSize, total - p);

			byte[] frag = new byte[size];
			System.arraycopy(payload, p, frag, 0, frag.length);

			// filling header
			byte[] header = new byte[13];
			header[0] = HEADER;
			Util.intToByteArray(header, 1, messageId);
			Util.intToByteArray(header, 5, total);
			Util.intToByteArray(header, 9, p);
			
			if (logger.isLoggable(Level.FINEST)) {
				String s = "generating frag: id({0}) total({1}) offset({2})";
				logger.log(Level.FINEST, s, new Object[] { messageId,
						total, p });
			}

			G2CLMessage fragMsg = (G2CLMessage) data.createMessage();
			fragMsg.setSenderAddress(message.getSenderAddress());
			fragMsg.setPayload(frag);
			fragMsg.addHeader(header);
			p += frag.length;
			return fragMsg;
		}
	}
}

/** Class to help on re-assembling of messages */
class AsmMessage {
	private static Logger logger = Logger.getLogger(AsmMessage.class
			.getName());
	private byte[] message;
	
	private int countBytes;
	
	public AsmMessage(int size) {
		message = new byte[size];
		countBytes = 0;
	}
	
	public void add(byte[] fragment, int offset) {
		if (logger.isLoggable(Level.FINEST)) {
			logger.finest("==============\n" +
			"original message size: " + message.length + "\n" +
			"bytes received: " + countBytes + "\n" +
			"frag added size: " + fragment.length + "\n" +
			"frag offset: " + offset + "\n" +
			"================");
		}
		countBytes += fragment.length;
		System.arraycopy(fragment, 0, message, offset, fragment.length);
	}
	
	public boolean isComplete() {
		return countBytes == message.length;
	}
	
	public byte[] getMessage() {
		return message;
	}
}
